// Copyright (C) 2010  Andrew H. Chan
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
// 
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// The MultiM class is the data structure that contains the sample configurations (set of haplotypes).
#ifndef _MULTI_M_H__
#define _MULTI_M_H__

#include "basic_functions.h"

class MultiM {
    public:
        // Comparison functor for index_t to use with std::map
        class index_cmp_t : public std::binary_function<index_t, index_t, bool> {
            public:
                // predicate: a < b
                bool operator()(index_t const& a, index_t const& b) const {
                    if(a.size() != b.size()) {
                        return a.size() < b.size();
                    }
                    for(index_t::size_type i = 0; i < a.size(); ++i) {
                        if(a[i]!=b[i]) {
                            return a[i]<b[i];
                        }
                    }
                    return false;
                }
        };

        // Use a hash table instead of map for better efficiency.
        // I'm using a map right now because I need to be able to hash a sample configuration, which is of data type MultiM. std::map's have a well-defined ordering on its elements (as long as every key is unique). An unordered_map is not guaranteed to have any ordering on its elements. Thus using std::map allows me to hash a sample configuration whereas an unordered map cannot be hashed reliably.
        typedef std::map<index_t, int, index_cmp_t> data_type; 
        //typedef boost::unordered_map<index_t, int> data_type; 
        typedef data_type::iterator iterator; 
        typedef data_type::const_iterator const_iterator; 

        class ProxyRef;

        // data members
        data_type data;

        // constructor
        MultiM () {}

        // functions
        size_t size() const {
            size_t tot_size = 0;
            // Iterate through index vector and add up entries
            for(data_type::const_iterator iter = data.begin();
                    iter != data.end(); ++iter) {
                tot_size += iter->second;
            }
            return tot_size;
        }

        void clear() {
            data.clear();
        }

        size_t distinct_hap() const {
            // Counts the number of distinct haplotypes (number of indexes stored in data)
            return data.size();
        }

        // Count the number of loci defined by sample
        // NOTE: This is weighted by quantity
        // This is what is referred to by "degree" in Jenkins, Song (2009)
        // My use of the word degree in the rest of the program is different
        size_t compute_V() const {
            size_t V = 0;
            // Go through all indexes
            for(data_type::const_iterator outer_iter = data.begin();
                    outer_iter != data.end(); ++outer_iter) {
                // Count number of loci defined by current index
                size_t num_def_loci = 0;
                for(index_t::const_iterator inner_iter = 
                        outer_iter->first.begin();
                        inner_iter != outer_iter->first.end();
                        ++inner_iter) {
                    if(*inner_iter != -1) {
                        ++num_def_loci;
                    }
                }
                size_t quantity = outer_iter->second;
                V += num_def_loci * quantity; 
            }
            return V;
        }

        size_t compute_U() const {
            size_t U = 0;
            // Go through all indexes (individuals in sample)
            for(data_type::const_iterator outer_iter = data.begin();
                    outer_iter != data.end(); ++outer_iter) {
                // Count number of significant breakpoints in current index
                // (i.e. the breakpoints we care about)
                // Equal to right_endpoint - left_endpoint
                size_t left_point=0;
                size_t right_point=0;
                index_t const &index = outer_iter->first;
                // get left endpoint
                for(uint ii=0; ii<index.size(); ++ii) {
                    if(index[ii] != -1) {
                        left_point = ii;
                        break;
                    }
                }
                // get right endpoint
                for(uint ii=index.size()-1; ii>=0; --ii) {
                    if(index[ii] != -1) {
                        right_point = ii;
                        break;
                    }
                }
                
                int num_breakpoints = right_point - left_point;
                assert(num_breakpoints >= 0);

                size_t quantity = outer_iter->second;
                U += num_breakpoints * quantity; 
            }
            return U;
        }

        iterator begin() {
            return data.begin();
        }

        const_iterator begin() const {
            return data.begin();
        }

        iterator end() {
            return data.end();
        }

        const_iterator end() const {
            return data.end();
        }

        bool empty() const {
            return data.empty();
        }

        // operators
        ProxyRef operator[](index_t const& index) {
            return ProxyRef(index, data);
        }

        bool operator==(MultiM const& other) const {
            return data == other.data;
        }

        bool operator!=(MultiM const& other) const {
            return !(data == other.data);
        }

        class ProxyRef {
            private:
                index_t index;
                data_type &data;
            public:
                ProxyRef(index_t const& index, data_type &data) : index(index), data(data) {}

                // Conversion to int 
                // If a coordinate is not mapped, its value is by default 0
                operator int() const {
                    if(data.find(index) == data.end()) {
                        return 0;
                    }
                    else {
                        return data[index];
                    }
                }

                ProxyRef const& operator=(int value) {
                    if(value == 0) {
                        // remove from map. Note, erase() won't throw an exception if index isn't in map.
                        data.erase(index);
                    }
                    else {
                        data[index] = value;
                    }
                    return *this;
                }

                ProxyRef const& operator=(ProxyRef const& ref) {
                    if(this == &ref) {
                        return *this;
                    }
                    if(ref == 0) {
                        // remove from map. Note, erase() won't throw an exception if index isn't in map.
                        data.erase(index);
                    }
                    else {
                        data[index] = ref;
                    }
                    return *this;
                }

                ProxyRef const& operator++() {
                    ++data[index];
                    return *this;
                }

                data_type::mapped_type const operator++(int) {
                    data_type::mapped_type tmp_value = data[index];
                    ++data[index];
                    return tmp_value;
                }

                ProxyRef const& operator--() {
                    --data[index];
                    if(data[index] == 0) {
                        // remove from map 
                        data.erase(index);
                    }
                    return *this;
                }

                data_type::mapped_type const operator--(int) {
                    data_type::mapped_type tmp_value = data[index];
                    --data[index];
                    if(data[index] == 0) {
                        // remove from map 
                        data.erase(index);
                    }
                    return tmp_value;
                }
        };
};

// hash function for MultiM
size_t hash_value(MultiM const& multiM);

#endif
